import sys
import os
try:
    parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    if parent_dir not in sys.path:
        sys.path.append(parent_dir)
    from base_code.common import *
    from base_code import *  # load pocketflow
except ImportError as e:
    print(f"❌ 错误: 无法导入 'base_code.common'。请确保文件结构正确。")
    print(f"   错误详情: {e}")


import functools
import asyncio
import ollama
import re
import json
from typing import List, Dict, Any, Optional, Callable

class CommandRegistry(ToolSystem):
    def __init__(self):
        super().__init__()
        self.lines = [
            "     欢迎使用 ReAct plus v0.3 (工作流版)",
            " ",
            "支持自然语言交互，所有逻辑由Agent的P-D-A循环由工作流处理",
            " ",
            "use_tool <tool> <param>  - 使用工具",
            "ls path                  - 列出path目录中的文件列表",
            "readpdf <pdf>            - 列出path目录中的文件列表",
            "help                     - 显示帮助信息",
            "exit                     - 退出程序"
        ]

    def get_tools_description(self) -> str:
        """为 LLM 生成格式化的工具描述列表。"""
        if not self._tools:
            return "  - 无可用工具。"
        lines = []
        for name, tool in self._tools.items():
            # 为 'finish' 等内置控制工具提供特殊描述
            if name == 'finish':
                lines.append(
                    "  - finish(answer: str): 当你完成所有任务并得出最终结论时调用，在 'answer' 中提供最终答案。")
                continue

            args_list = list(tool.required_args) + [f"{k}={v}" for k, v in (tool.optional_args or {}).items()]
            # 过滤掉内部参数 'memory'
            args_list = [arg for arg in args_list if arg != 'memory']
            args_str = ", ".join(args_list)
            lines.append(f"  - {tool.name}({args_str}): {tool.description}")
        return "\n".join(lines)


##################################################################
# 全局单例对象：工具系统（环境层）
##################################################################
GLOBAL_TOOL_SYSTEM = CommandRegistry()


# ==============================================================================
#  使用装饰器注册一个本地工具作为示例
# ==============================================================================
@tool(tool_system=GLOBAL_TOOL_SYSTEM, name="add", description="计算两个数的和", required_args=["a", "b"])
def local_add_tool(a: int, b: int) -> int:
    """一个简单的本地加法工具。"""
    return a + b


##################################################################
# 基于LLM的动态工具测试：生成工具列表，让LLM生成调用
##################################################################
class FuncTestAgent:
    def __init__(self, registry, model_name="qwen3:4b"):
        self.registry = registry
        self.model_name = model_name
        self.ollama_client = ollama.AsyncClient()

    def _generate_prompt(self):
        system_prompt = """
        You are a ReAct-based testing assistant that helps automatically generate JSON test cases 
        for a list of registered Python tools.
    
        Each tool is defined with its name, parameters (with optional defaults), and description. 
        Your job is to generate a JSON array where each element represents one tool call test case.
    
        Each JSON element must include:
        - "name": the tool's function name
        - "arguments": a dictionary of arguments to test this function
          - required parameters must always have a value
          - optional parameters (with defaults) may be included or omitted
    
        Rules:
        - Generate realistic placeholder values (e.g., strings, numbers, file paths)
        - Use simple examples that can execute safely in a local environment
        - The output must be valid JSON that can be parsed by `json.loads()`
        - Do NOT include Markdown formatting, comments, or extra text — JSON only
    
        Example input:
        =====================
        工具列表:
        - add(a, b): 计算两个数的和
        - fibonacci(n): 计算前 n 项斐波那契数列
        - weather(location): 简单天气工具示例
        - read_pdf(path, max_pages=1): 读取 PDF 文件前 max_pages 页
        - list_files(directory=.): 列出指定目录
        - system_info(): 返回系统信息
        =====================
    
        Output only the final JSON result.
        Example valid output:
        [
          {"name": "add", "arguments": {"a": 5, "b": 10}},
          {"name": "fibonacci", "arguments": {"n": 7}},
          {"name": "weather", "arguments": {"location": "Beijing"}},
          {"name": "read_pdf", "arguments": {"path": "example.pdf", "max_pages": 1}},
          {"name": "list_files", "arguments": {}},
          {"name": "system_info", "arguments": {}}
        ]
        """
        return system_prompt

    async def generate_test_cases(self):
        system_prompt = self._generate_prompt()
        get_tools_description = self.registry.get_tools_description()
        user_prompt = f"Here is the tool list:\n{get_tools_description}\nPlease generate JSON test cases."
        response = await self.ollama_client.chat(
            model=self.model_name,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ]#,
            #format="json"
        )

        content = response["message"]["content"]
        print(content)

        def clean_think_blocks(text: str) -> str:
            """去除 <think>...</think> 推理块"""
            return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()

        content = clean_think_blocks(content)
        if content.startswith("```"):
            content = content.split("```")[1].replace("json", "").strip()
        try:
            llm_output = json.loads(content)
            return llm_output
        except Exception as e:
            print("⚠️ JSON解析失败，返回原始内容：", content)
            raise e

    async def run_tests(self):
        json_list = await self.generate_test_cases()

        for entry in json_list:
            name = entry["name"]
            args = entry.get("arguments", {})
            tool_func = self.registry.get_tool(name)
            if not tool_func:
                print(f"❌ 工具 {name} 未注册。")
                continue

            print(f"🧩 测试 {name}({args}) ...")
            try:
                result = tool_func(**args)
                if asyncio.iscoroutine(result):
                    result = await result
                print(f"✅ {name} 返回: {result}\n")
            except Exception as e:
                print(f"❌ {name} 执行错误: {e}\n")


# ==============================================================================
# 2. HelpAgent 支持UI绘制
# ==============================================================================
class HelpAgent:
    def __init__(self, registry: CommandRegistry):
        self.command_registry = registry

    def show_intro(self):
        try:
            return create_aligned_intro(self.command_registry.lines)
        except Exception:
            return "\n".join(self.command_registry.lines)

    def show_help(self, command=None):
        if command and self.command_registry.has_tool(command):
            t = self.command_registry.get_tool(command)
            usage = f"{t.name} {' '.join(t.required_args)}"
            return f"命令: {t.name}\n功能: {t.description}\n用法: {usage}"
        else:
            return self.show_intro()

async def main():
    await GLOBAL_TOOL_SYSTEM.load_tools_from_mcp("http://127.0.0.1:9000")
    print("\n===================== 🧰 工具列表 =====================")
    print(GLOBAL_TOOL_SYSTEM.get_tools_description())
    print("======================================================\n")

    # 3️⃣ 自动测试每个工具调用
    tools = GLOBAL_TOOL_SYSTEM.get_tools()
    print(f"共加载 {len(tools)} 个工具，开始测试调用...\n")

    agent = FuncTestAgent(GLOBAL_TOOL_SYSTEM)

    # await agent.run_tests()

    agent = HelpAgent(GLOBAL_TOOL_SYSTEM)
    print(agent.show_intro())
    agent.show_help()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n程序被用户中断。")


